import pandas as pd
import numpy as np
import matplotlib.pyplot as plt


def load_metrix(name, indicator='val_mean_squared_error'):
    '''
    :param indicator: val_loss val_mean_squared_error  val_r_2_score r_2_score
    :param name: gru,lstm,saes
    :return:
    '''
    df = pd.read_csv('model/' + name + ' loss.csv', encoding='utf-8')
    metrix = df[indicator]
    return metrix


def plot_mae():
    names = ['gru', 'saes', 'lstm']
    for name in names:
        mae = load_metrix(name)
        plt.plot(np.arange(0, len(mae)), mae, label=name)
    plt.xlabel('epoch')
    plt.ylabel('Mean Average Error')
    plt.grid(True)
    plt.legend()
    plt.savefig('images/MAE versus Epoch')
    plt.show()


def plot_r2():
    names = ['gru', 'saes', 'lstm']
    for name in names:
        score = load_metrix(name,indicator='r_2_score')
        plt.plot(np.arange(0, len(score)), score, label=name)
    plt.xlabel('epoch')
    plt.ylabel('R2 Score')
    plt.grid(True)
    plt.legend()
    plt.savefig('images/R2Score versus Epoch')
    plt.show()


if __name__ == '__main__':
    # plot_mae()
    plot_r2()
